// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/functions/prestosql/comparisons.h>
#include <pollux/common/base/exceptions.h>
#include <pollux/functions/udf.h>
#include <pollux/functions/lib/simd_comparison_util.h>
#include <pollux/vector/base_vector.h>

namespace kumo::pollux::functions {
    using Eq = std::equal_to<>;
    using Neq = std::not_equal_to<>;
    using Lt = std::less<>;
    using Lte = std::less_equal<>;
    using Gt = std::greater<>;
    using Gte = std::greater_equal<>;

    namespace {
        /// This class implements comparison for vectors of primitive types using SIMD.
        /// Currently this only supports fixed length primitive types (except Boolean).
        /// It also requires the vectors to have a flat encoding.
        /// If the vector encoding is not flat, we revert to non simd approach.
        template<typename ComparisonOp, typename Arch = ksimd::default_arch>
        struct SimdComparator {
            template<typename T>
            inline bool compare(T &l, T &r) const {
                if constexpr (std::is_floating_point_v<T>) {
                    bool filtered = false;
                    if constexpr (std::is_same_v<ComparisonOp, Eq>) {
                        filtered = util::floating_point::NaNAwareEquals<T>{}(l, r);
                    } else if constexpr (std::is_same_v<ComparisonOp, Neq>) {
                        filtered = !util::floating_point::NaNAwareEquals<T>{}(l, r);
                    } else if constexpr (std::is_same_v<ComparisonOp, Lt>) {
                        filtered = util::floating_point::NaNAwareLessThan<T>{}(l, r);
                    } else if constexpr (std::is_same_v<ComparisonOp, Lte>) {
                        filtered = util::floating_point::NaNAwareLessThanEqual<T>{}(l, r);
                    } else if constexpr (std::is_same_v<ComparisonOp, Gt>) {
                        filtered = util::floating_point::NaNAwareGreaterThan<T>{}(l, r);
                    } else if constexpr (std::is_same_v<ComparisonOp, Gte>) {
                        filtered = util::floating_point::NaNAwareGreaterThanEqual<T>{}(l, r);
                    }
                    return filtered;
                } else {
                    return ComparisonOp()(l, r);
                }
            }

            template<
                TypeKind kind,
                typename std::enable_if_t<
                    (ksimd::has_simd_register<
                         typename TypeTraits<kind>::NativeType>::value &&
                     kind != TypeKind::BOOLEAN) ||
                    kind == TypeKind::HUGEINT,
                    int>  = 0>
            void applyComparison(
                const SelectivityVector &rows,
                BaseVector &lhs,
                BaseVector &rhs,
                exec::EvalCtx &context,
                VectorPtr &result) {
                using T = typename TypeTraits<kind>::NativeType;

                auto resultVector = result->as_unchecked<FlatVector<bool> >();
                auto rawResult = resultVector->mutableRawValues<uint8_t>();

                bool isSimdizable = (lhs.is_constant_encoding() || lhs.is_flat_encoding()) &&
                                    (rhs.is_constant_encoding() || rhs.is_flat_encoding()) &&
                                    rows.isAllSelected();

                static const bool isTypeNotSupported =
                        std::is_same_v<T, int128_t> || std::is_floating_point_v<T>;

                if (!isSimdizable || isTypeNotSupported) {
                    exec::LocalDecodedVector lhsDecoded(context, lhs, rows);
                    exec::LocalDecodedVector rhsDecoded(context, rhs, rows);

                    context.applyToSelectedNoThrow(rows, [&](auto row) {
                        auto l = lhsDecoded->template value_at<T>(row);
                        auto r = rhsDecoded->template value_at<T>(row);
                        auto filtered = compare(l, r);
                        resultVector->set(row, filtered);
                    });
                    return;
                }

                if constexpr (!isTypeNotSupported) {
                    if (lhs.is_constant_encoding() && rhs.is_constant_encoding()) {
                        auto l = lhs.as_unchecked<ConstantVector<T> >()->value_at(0);
                        auto r = rhs.as_unchecked<ConstantVector<T> >()->value_at(0);
                        applySimdComparison<T, true, true, ComparisonOp>(
                            rows.begin(), rows.end(), &l, &r, rawResult);
                    } else if (lhs.is_constant_encoding()) {
                        auto l = lhs.as_unchecked<ConstantVector<T> >()->value_at(0);
                        auto rawRhs = rhs.as_unchecked<FlatVector<T> >()->rawValues();
                        applySimdComparison<T, true, false, ComparisonOp>(
                            rows.begin(), rows.end(), &l, rawRhs, rawResult);
                    } else if (rhs.is_constant_encoding()) {
                        auto rawLhs = lhs.as_unchecked<FlatVector<T> >()->rawValues();
                        auto r = rhs.as_unchecked<ConstantVector<T> >()->value_at(0);
                        applySimdComparison<T, false, true, ComparisonOp>(
                            rows.begin(), rows.end(), rawLhs, &r, rawResult);
                    } else {
                        auto rawLhs = lhs.as_unchecked<FlatVector<T> >()->rawValues();
                        auto rawRhs = rhs.as_unchecked<FlatVector<T> >()->rawValues();
                        applySimdComparison<T, false, false, ComparisonOp>(
                            rows.begin(), rows.end(), rawLhs, rawRhs, rawResult);
                    }

                    resultVector->clear_nulls(rows);
                }
            }

            template<
                TypeKind kind,
                typename std::enable_if_t<
                    (!ksimd::has_simd_register<
                         typename TypeTraits<kind>::NativeType>::value ||
                     kind == TypeKind::BOOLEAN) &&
                    kind != TypeKind::HUGEINT,
                    int>  = 0>
            void applyComparison(
                const SelectivityVector & /* rows */,
                BaseVector & /* lhs */,
                BaseVector & /* rhs */,
                exec::EvalCtx & /* context */,
                VectorPtr & /* result */) {
                POLLUX_UNSUPPORTED("Unsupported type for SIMD comparison");
            }
        };

        template<typename ComparisonOp, typename Arch = ksimd::default_arch>
        class ComparisonSimdFunction : public exec::VectorFunction {
        public:
            void apply(
                const SelectivityVector &rows,
                std::vector<VectorPtr> &args,
                const TypePtr &outputType,
                exec::EvalCtx &context,
                VectorPtr &result) const override {
                POLLUX_CHECK_EQ(args.size(), 2, "Comparison requires two arguments");
                POLLUX_CHECK_EQ(args[0]->type_kind(), args[1]->type_kind());
                POLLUX_USER_CHECK_EQ(outputType->kind(), TypeKind::BOOLEAN);

                context.ensure_writable(rows, outputType, result);
                auto comparator = SimdComparator<ComparisonOp>{};

                if (args[0]->type()->isLongDecimal()) {
                    comparator.template applyComparison<TypeKind::HUGEINT>(
                        rows, *args[0], *args[1], context, result);
                    return;
                }

                POLLUX_DYNAMIC_SCALAR_TYPE_DISPATCH(
                    comparator.template applyComparison,
                    args[0]->type_kind(),
                    rows,
                    *args[0],
                    *args[1],
                    context,
                    result);
            }

            static std::vector<std::shared_ptr<exec::FunctionSignature> > signatures() {
                std::vector<std::shared_ptr<exec::FunctionSignature> > signatures;

                for (const auto &inputType: {
                         "tinyint",
                         "smallint",
                         "integer",
                         "bigint",
                         "real",
                         "double",
                         "date",
                         "interval day to second",
                         "interval year to month",
                     }) {
                    signatures.push_back(exec::FunctionSignatureBuilder()
                        .returnType("boolean")
                        .argumentType(inputType)
                        .argumentType(inputType)
                        .build());
                }
                signatures.push_back(exec::FunctionSignatureBuilder()
                    .integerVariable("a_precision")
                    .integerVariable("a_scale")
                    .returnType("boolean")
                    .argumentType("DECIMAL(a_precision, a_scale)")
                    .argumentType("DECIMAL(a_precision, a_scale)")
                    .build());
                return signatures;
            }

            bool supportsFlatNoNullsFastPath() const override {
                return true;
            }

            exec::FunctionCanonicalName getCanonicalName() const override {
                return std::is_same_v<ComparisonOp, Lt>
                           ? exec::FunctionCanonicalName::kLt
                           : exec::FunctionCanonicalName::kUnknown;
            }
        };
    } // namespace

    POLLUX_DECLARE_VECTOR_FUNCTION(
        udf_simd_comparison_eq,
        (ComparisonSimdFunction<Eq>::signatures()),
        (std::make_unique<ComparisonSimdFunction<Eq>>()));

    POLLUX_DECLARE_VECTOR_FUNCTION(
        udf_simd_comparison_neq,
        (ComparisonSimdFunction<Neq>::signatures()),
        (std::make_unique<ComparisonSimdFunction<Neq>>()));

    POLLUX_DECLARE_VECTOR_FUNCTION(
        udf_simd_comparison_lt,
        (ComparisonSimdFunction<Lt>::signatures()),
        (std::make_unique<ComparisonSimdFunction<Lt>>()));

    POLLUX_DECLARE_VECTOR_FUNCTION(
        udf_simd_comparison_gt,
        (ComparisonSimdFunction<Gt>::signatures()),
        (std::make_unique<ComparisonSimdFunction<Gt>>()));

    POLLUX_DECLARE_VECTOR_FUNCTION(
        udf_simd_comparison_lte,
        (ComparisonSimdFunction<Lte>::signatures()),
        (std::make_unique<ComparisonSimdFunction<Lte>>()));

    POLLUX_DECLARE_VECTOR_FUNCTION(
        udf_simd_comparison_gte,
        (ComparisonSimdFunction<Gte>::signatures()),
        (std::make_unique<ComparisonSimdFunction<Gte>>()));
} // namespace kumo::pollux::functions
